// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/array/array_base.h>
#include <nebula/compute/api_scalar.h>
#include <nebula/compute/cast.h>
#include <nebula/compute/kernels/common_internal.h>
#include <nebula/compute/kernels/util_internal.h>
#include <nebula/types/type.h>
#include <nebula/bits/bit_util.h>
#include <nebula/bits/bitmap_writer.h>
#include <nebula/util/hashing.h>
#include <nebula/core/visit_data_inline.h>

namespace nebula {

    using internal::HashTraits;
}
namespace nebula::compute::internal {
    namespace {

        // This base class enables non-templated access to the value set type
        struct SetLookupStateBase : public KernelState {
            std::shared_ptr<DataType> value_set_type;
        };

        template<typename Type>
        struct SetLookupState : public SetLookupStateBase {
            explicit SetLookupState(MemoryPool *pool) : memory_pool(pool) {}

            turbo::Status init(const SetLookupOptions &options) {
                this->null_matching_behavior = options.GetNullMatchingBehavior();
                if (options.value_set.is_array()) {
                    const ArrayData &value_set = *options.value_set.array();
                    memo_index_to_value_index.reserve(value_set.length);
                    lookup_table =
                            MemoTable(memory_pool,
                                      ::nebula::internal::HashTable<char>::kLoadFactor * value_set.length);
                    TURBO_RETURN_NOT_OK(AddArrayValueSet(options, *options.value_set.array()));
                } else if (options.value_set.kind() == Datum::CHUNKED_ARRAY) {
                    const ChunkedArray &value_set = *options.value_set.chunked_array();
                    memo_index_to_value_index.reserve(value_set.length());
                    lookup_table =
                            MemoTable(memory_pool,
                                      ::nebula::internal::HashTable<char>::kLoadFactor * value_set.length());

                    int64_t offset = 0;
                    for (const std::shared_ptr<Array> &chunk: value_set.chunks()) {
                        TURBO_RETURN_NOT_OK(AddArrayValueSet(options, *chunk->data(), offset));
                        offset += chunk->length();
                    }
                } else {
                    return turbo::invalid_argument_error("value_set should be an array or chunked array");
                }
                if (this->null_matching_behavior != SetLookupOptions::SKIP &&
                    lookup_table->GetNull() >= 0) {
                    null_index = memo_index_to_value_index[lookup_table->GetNull()];
                }
                value_set_type = options.value_set.type();
                return turbo::OkStatus();
            }

            turbo::Status AddArrayValueSet(const SetLookupOptions &options, const ArrayData &data,
                                           int64_t start_index = 0) {
                using T = typename GetViewType<Type>::T;
                int32_t index = static_cast<int32_t>(start_index);
                auto visit_valid = [&](T v) {
                    const auto memo_size = static_cast<int32_t>(memo_index_to_value_index.size());
                    int32_t unused_memo_index;
                    // (capture `memo_size` by value because of ARROW-17567)
                    auto on_found = [&, memo_size](int32_t memo_index) {
                        DKCHECK_LT(memo_index, memo_size);
                    };
                    auto on_not_found = [&, memo_size](int32_t memo_index) {
                        DKCHECK_EQ(memo_index, memo_size);
                        memo_index_to_value_index.push_back(index);
                    };
                    TURBO_RETURN_NOT_OK(lookup_table->GetOrInsert(
                            v, std::move(on_found), std::move(on_not_found), &unused_memo_index));
                    ++index;
                    return turbo::OkStatus();
                };
                auto visit_null = [&]() {
                    const auto memo_size = static_cast<int32_t>(memo_index_to_value_index.size());
                    auto on_found = [&, memo_size](int32_t memo_index) {
                        DKCHECK_LT(memo_index, memo_size);
                    };
                    auto on_not_found = [&, memo_size](int32_t memo_index) {
                        DKCHECK_EQ(memo_index, memo_size);
                        memo_index_to_value_index.push_back(index);
                    };
                    lookup_table->GetOrInsertNull(std::move(on_found), std::move(on_not_found));
                    ++index;
                    return turbo::OkStatus();
                };

                return VisitArraySpanInline<Type>(data, visit_valid, visit_null);
            }

            using MemoTable = typename HashTraits<Type>::MemoTableType;
            std::optional<MemoTable> lookup_table;  // use optional for delayed initialization
            MemoryPool *memory_pool;
            // When there are duplicates in value_set, the MemoTable indices must
            // be mapped back to indices in the value_set.
            std::vector<int32_t> memo_index_to_value_index;
            int32_t null_index = -1;
            SetLookupOptions::NullMatchingBehavior null_matching_behavior;
        };

        template<>
        struct SetLookupState<NullType> : public SetLookupStateBase {
            explicit SetLookupState(MemoryPool *) {}

            turbo::Status init(SetLookupOptions &options) {
                null_matching_behavior = options.GetNullMatchingBehavior();
                value_set_has_null = (options.value_set.length() > 0) &&
                                     this->null_matching_behavior != SetLookupOptions::SKIP;
                value_set_type = null();
                return turbo::OkStatus();
            }

            bool value_set_has_null;
            SetLookupOptions::NullMatchingBehavior null_matching_behavior;
        };

// TODO: Put this concept somewhere reusable
        template<int width>
        struct UnsignedIntType;

        template<>
        struct UnsignedIntType<1> {
            using Type = UInt8Type;
        };

        template<>
        struct UnsignedIntType<2> {
            using Type = UInt16Type;
        };

        template<>
        struct UnsignedIntType<4> {
            using Type = UInt32Type;
        };

        template<>
        struct UnsignedIntType<8> {
            using Type = UInt64Type;
        };

// Constructing the type requires a type parameter
        struct InitStateVisitor {
            KernelContext *ctx;
            SetLookupOptions options;
            TypeHolder arg_type;
            std::unique_ptr<KernelState> result;

            InitStateVisitor(KernelContext *ctx, const KernelInitArgs &args)
                    : ctx(ctx),
                      options(*turbo::checked_cast<const SetLookupOptions *>(args.options)),
                      arg_type(args.inputs[0]) {}

            template<typename Type>
            turbo::Status init() {
                using StateType = SetLookupState<Type>;
                result.reset(new StateType(ctx->exec_context()->memory_pool()));
                return static_cast<StateType *>(result.get())->init(options);
            }

            turbo::Status Visit(const DataType &) { return init<NullType>(); }

            template<typename Type>
            enable_if_boolean<Type, turbo::Status> Visit(const Type &) {
                return init<BooleanType>();
            }

            template<typename Type>
            enable_if_t<has_c_type<Type>::value && !is_boolean_type<Type>::value &&
                        !std::is_same<Type, MonthDayNanoIntervalType>::value,
                    turbo::Status>
            Visit(const Type &) {
                return init<typename UnsignedIntType<sizeof(typename Type::c_type)>::Type>();
            }

            template<typename Type>
            enable_if_base_binary<Type, turbo::Status> Visit(const Type &) {
                return init<typename Type::PhysicalType>();
            }

            // Handle Decimal128Type, FixedSizeBinaryType
            turbo::Status Visit(const FixedSizeBinaryType &type) { return init<FixedSizeBinaryType>(); }

            turbo::Status Visit(const MonthDayNanoIntervalType &type) {
                return init<MonthDayNanoIntervalType>();
            }

            turbo::Result<std::unique_ptr<KernelState>> get_result() {
                if (arg_type.id() == Type::TIMESTAMP &&
                    options.value_set.type()->id() == Type::TIMESTAMP) {
                    // Other types will fail when casting, so no separate check is needed
                    const auto &ty1 = turbo::checked_cast<const TimestampType &>(*arg_type);
                    const auto &ty2 = turbo::checked_cast<const TimestampType &>(*options.value_set.type());
                    if (ty1.timezone().empty() ^ ty2.timezone().empty()) {
                        return turbo::failed_precondition_error(
                                "Cannot compare timestamp with timezone to timestamp without timezone, got: ",
                                ty1, " and ", ty2);
                    }
                } else if ((arg_type.id() == Type::STRING || arg_type.id() == Type::LARGE_STRING) &&
                           !is_base_binary_like(options.value_set.type()->id())) {
                    // This is a bit of a hack, but don't implicitly cast from a non-binary
                    // type to string, since most types support casting to string and that
                    // may lead to surprises. However, we do want most other implicit casts.
                    return turbo::failed_precondition_error("Array type doesn't match type of values set: ",
                                                            *arg_type,
                                                            " vs ", *options.value_set.type());
                }

                if (!options.value_set.is_arraylike()) {
                    return turbo::invalid_argument_error("Set lookup value set must be Array or ChunkedArray");
                } else if (!options.value_set.type()->equals(*arg_type)) {
                    auto cast_result =
                            Cast(options.value_set, CastOptions::Safe(arg_type.get_shared_ptr()),
                                 ctx->exec_context());
                    if (cast_result.ok()) {
                        options.value_set = *cast_result;
                    } else if (CanCast(*arg_type.type, *options.value_set.type())) {
                        // Avoid casting from non binary types to string like above
                        // Otherwise, will try to cast input array to value set type during kernel exec
                        if ((options.value_set.type()->id() == Type::STRING ||
                             options.value_set.type()->id() == Type::LARGE_STRING) &&
                            !is_base_binary_like(arg_type.id())) {
                            return turbo::failed_precondition_error("Array type doesn't match type of values set: ",
                                                                    *arg_type, " vs ", *options.value_set.type());
                        }
                    } else {
                        return turbo::failed_precondition_error("Array type doesn't match type of values set: ",
                                                                *arg_type, " vs ", *options.value_set.type());
                    }
                }

                TURBO_RETURN_NOT_OK(visit_type_inline(*options.value_set.type(), this));
                return std::move(result);
            }
        };

        turbo::Result<std::unique_ptr<KernelState>> InitSetLookup(KernelContext *ctx,
                                                                  const KernelInitArgs &args) {
            if (args.options == nullptr) {
                return turbo::invalid_argument_error(
                        "Attempted to call a set lookup function without SetLookupOptions");
            }

            return InitStateVisitor{ctx, args}.get_result();
        }

        struct IndexInVisitor {
            KernelContext *ctx;
            const ArraySpan &data;
            ArraySpan *out;
            uint8_t *out_bitmap;

            IndexInVisitor(KernelContext *ctx, const ArraySpan &data, ArraySpan *out)
                    : ctx(ctx), data(data), out(out), out_bitmap(out->buffers[0].data) {}

            turbo::Status Visit(const DataType &type) {
                DKCHECK(false) << "IndexIn " << type;
                return turbo::unimplemented_error("IndexIn has no implementation with value type ", type);
            }

            turbo::Status Visit(const NullType &) {
                const auto &state = turbo::checked_cast<const SetLookupState<NullType> &>(*ctx->state());

                if (data.length != 0) {
                    bit_util::SetBitsTo(out_bitmap, out->offset, out->length,
                                        state.null_matching_behavior == SetLookupOptions::MATCH &&
                                        state.value_set_has_null);

                    // Set all values to 0, which will be unmasked only if null is in the value_set
                    // and null_matching_behavior is equal to MATCH
                    std::memset(out->get_values<int32_t>(1), 0x00, out->length * sizeof(int32_t));
                }
                return turbo::OkStatus();
            }

            template<typename Type>
            turbo::Status ProcessIndexIn(const SetLookupState<Type> &state, const ArraySpan &input) {
                using T = typename GetViewType<Type>::T;
                FirstTimeBitmapWriter bitmap_writer(out_bitmap, out->offset, out->length);
                int32_t *out_data = out->get_values<int32_t>(1);
                VisitArraySpanInline<Type>(
                        input,
                        [&](T v) {
                            int32_t index = state.lookup_table->Get(v);
                            if (index != -1) {
                                bitmap_writer.Set();

                                // matching needle; output index from value_set
                                *out_data++ = state.memo_index_to_value_index[index];
                            } else {
                                // no matching needle; output null
                                bitmap_writer.Clear();
                                *out_data++ = 0;
                            }
                            bitmap_writer.next();
                        },
                        [&]() {
                            if (state.null_index != -1 &&
                                state.null_matching_behavior == SetLookupOptions::MATCH) {
                                bitmap_writer.Set();

                                // value_set included null
                                *out_data++ = state.null_index;
                            } else {
                                // value_set does not include null; output null
                                bitmap_writer.Clear();
                                *out_data++ = 0;
                            }
                            bitmap_writer.next();
                        });
                bitmap_writer.finish();
                return turbo::OkStatus();
            }

            template<typename Type>
            turbo::Status ProcessIndexIn() {
                const auto &state = turbo::checked_cast<const SetLookupState<Type> &>(*ctx->state());
                if (!data.type->equals(state.value_set_type)) {
                    auto materialized_input = data.to_array_data();
                    auto cast_result = Cast(*materialized_input, state.value_set_type,
                                            CastOptions::Safe(), ctx->exec_context());
                    if (TURBO_UNLIKELY(!cast_result.ok())) {
                        if (turbo::is_unimplemented(cast_result.status())) {
                            return turbo::failed_precondition_error("Array type doesn't match type of values set: ",
                                                                    *data.type, " vs ", *state.value_set_type);
                        }
                        return cast_result.status();
                    }
                    auto casted_input = *cast_result;
                    return ProcessIndexIn(state, *casted_input.array());
                }
                return ProcessIndexIn(state, data);
            }

            template<typename Type>
            enable_if_boolean<Type, turbo::Status> Visit(const Type &) {
                return ProcessIndexIn<BooleanType>();
            }

            template<typename Type>
            enable_if_t<has_c_type<Type>::value && !is_boolean_type<Type>::value &&
                        !std::is_same<Type, MonthDayNanoIntervalType>::value,
                    turbo::Status>
            Visit(const Type &) {
                return ProcessIndexIn<
                        typename UnsignedIntType<sizeof(typename Type::c_type)>::Type>();
            }

            template<typename Type>
            enable_if_base_binary<Type, turbo::Status> Visit(const Type &) {
                return ProcessIndexIn<typename Type::PhysicalType>();
            }

            // Handle Decimal128Type, FixedSizeBinaryType
            turbo::Status Visit(const FixedSizeBinaryType &type) {
                return ProcessIndexIn<FixedSizeBinaryType>();
            }

            turbo::Status Visit(const MonthDayNanoIntervalType &type) {
                return ProcessIndexIn<MonthDayNanoIntervalType>();
            }

            turbo::Status execute() {
                const auto &state = turbo::checked_cast<const SetLookupStateBase &>(*ctx->state());
                return visit_type_inline(*state.value_set_type, this);
            }
        };

        turbo::Status ExecIndexIn(KernelContext *ctx, const ExecSpan &batch, ExecResult *out) {
            return IndexInVisitor(ctx, batch[0].array, out->array_span_mutable()).execute();
        }

// IsIn writes the results into a preallocated boolean data bitmap
        struct IsInVisitor {
            KernelContext *ctx;
            const ArraySpan &data;
            ArraySpan *out;
            uint8_t *out_boolean_bitmap;
            uint8_t *out_null_bitmap;

            IsInVisitor(KernelContext *ctx, const ArraySpan &data, ArraySpan *out)
                    : ctx(ctx),
                      data(data),
                      out(out),
                      out_boolean_bitmap(out->buffers[1].data),
                      out_null_bitmap(out->buffers[0].data) {}

            turbo::Status Visit(const DataType &type) {
                DKCHECK(false) << "IndexIn " << type;
                return turbo::unimplemented_error("IsIn has no implementation with value type ", type);
            }

            turbo::Status Visit(const NullType &) {
                const auto &state = turbo::checked_cast<const SetLookupState<NullType> &>(*ctx->state());

                if (state.null_matching_behavior == SetLookupOptions::MATCH &&
                    state.value_set_has_null) {
                    bit_util::SetBitsTo(out_boolean_bitmap, out->offset, out->length, true);
                    bit_util::SetBitsTo(out_null_bitmap, out->offset, out->length, true);
                } else if (state.null_matching_behavior == SetLookupOptions::SKIP ||
                           (!state.value_set_has_null &&
                            state.null_matching_behavior == SetLookupOptions::MATCH)) {
                    bit_util::SetBitsTo(out_boolean_bitmap, out->offset, out->length, false);
                    bit_util::SetBitsTo(out_null_bitmap, out->offset, out->length, true);
                } else {
                    bit_util::SetBitsTo(out_null_bitmap, out->offset, out->length, false);
                }
                return turbo::OkStatus();
            }

            template<typename Type>
            turbo::Status ProcessIsIn(const SetLookupState<Type> &state, const ArraySpan &input) {
                using T = typename GetViewType<Type>::T;
                FirstTimeBitmapWriter writer_boolean(out_boolean_bitmap, out->offset, out->length);
                FirstTimeBitmapWriter writer_null(out_null_bitmap, out->offset, out->length);
                bool value_set_has_null = state.null_index != -1;
                VisitArraySpanInline<Type>(
                        input,
                        [&](T v) {
                            if (state.lookup_table->Get(v) != -1) {  // true
                                writer_boolean.Set();
                                writer_null.Set();
                            } else if (state.null_matching_behavior == SetLookupOptions::INCONCLUSIVE &&
                                       value_set_has_null) {  // null
                                writer_boolean.Clear();
                                writer_null.Clear();
                            } else {  // false
                                writer_boolean.Clear();
                                writer_null.Set();
                            }
                            writer_boolean.next();
                            writer_null.next();
                        },
                        [&]() {
                            if (state.null_matching_behavior == SetLookupOptions::MATCH &&
                                value_set_has_null) {  // true
                                writer_boolean.Set();
                                writer_null.Set();
                            } else if (state.null_matching_behavior == SetLookupOptions::SKIP ||
                                       (!value_set_has_null && state.null_matching_behavior ==
                                                               SetLookupOptions::MATCH)) {  // false
                                writer_boolean.Clear();
                                writer_null.Set();
                            } else {  // null
                                writer_boolean.Clear();
                                writer_null.Clear();
                            }
                            writer_boolean.next();
                            writer_null.next();
                        });
                writer_boolean.finish();
                writer_null.finish();
                return turbo::OkStatus();
            }

            template<typename Type>
            turbo::Status ProcessIsIn() {
                const auto &state = turbo::checked_cast<const SetLookupState<Type> &>(*ctx->state());

                if (!data.type->equals(state.value_set_type)) {
                    auto materialized_input = data.to_array_data();
                    auto cast_result = Cast(*materialized_input, state.value_set_type,
                                            CastOptions::Safe(), ctx->exec_context());
                    if (TURBO_UNLIKELY(!cast_result.ok())) {
                        if (turbo::is_unimplemented(cast_result.status())) {
                            return turbo::failed_precondition_error("Array type doesn't match type of values set: ",
                                                                    *data.type, " vs ", *state.value_set_type);
                        }
                        return cast_result.status();
                    }
                    auto casted_input = *cast_result;
                    return ProcessIsIn(state, *casted_input.array());
                }
                return ProcessIsIn(state, data);
            }

            template<typename Type>
            enable_if_boolean<Type, turbo::Status> Visit(const Type &) {
                return ProcessIsIn<BooleanType>();
            }

            template<typename Type>
            enable_if_t<has_c_type<Type>::value && !is_boolean_type<Type>::value &&
                        !std::is_same<Type, MonthDayNanoIntervalType>::value,
                    turbo::Status>
            Visit(const Type &) {
                return ProcessIsIn<typename UnsignedIntType<sizeof(typename Type::c_type)>::Type>();
            }

            template<typename Type>
            enable_if_base_binary<Type, turbo::Status> Visit(const Type &) {
                return ProcessIsIn<typename Type::PhysicalType>();
            }

            // Handle Decimal128Type, FixedSizeBinaryType
            turbo::Status Visit(const FixedSizeBinaryType &type) {
                return ProcessIsIn<FixedSizeBinaryType>();
            }

            turbo::Status Visit(const MonthDayNanoIntervalType &type) {
                return ProcessIsIn<MonthDayNanoIntervalType>();
            }

            turbo::Status execute() {
                const auto &state = turbo::checked_cast<const SetLookupStateBase &>(*ctx->state());
                return visit_type_inline(*state.value_set_type, this);
            }
        };

        turbo::Status ExecIsIn(KernelContext *ctx, const ExecSpan &batch, ExecResult *out) {
            return IsInVisitor(ctx, batch[0].array, out->array_span_mutable()).execute();
        }

// Unary set lookup kernels available for the following input types
//
// * Null type
// * Boolean
// * Numeric
// * Simple temporal types (date, time, timestamp)
// * Base binary types
// * Decimal

        void AddBasicSetLookupKernels(ScalarKernel kernel,
                                      const std::shared_ptr<DataType> &out_ty,
                                      ScalarFunction *func) {
            auto AddKernels = [&](const std::vector<std::shared_ptr<DataType>> &types) {
                for (const std::shared_ptr<DataType> &ty: types) {
                    kernel.signature = KernelSignature::create({InputType(ty->id())}, out_ty);
                    KCHECK_OK(func->add_kernel(kernel));
                }
            };

            AddKernels(base_binary_types());
            AddKernels(numeric_types());
            AddKernels(temporal_types());
            AddKernels(duration_types());
            AddKernels({month_day_nano_interval()});

            std::vector<Type::type> other_types = {Type::BOOL, Type::DECIMAL128, Type::DECIMAL256,
                                                   Type::FIXED_SIZE_BINARY};
            for (auto ty: other_types) {
                kernel.signature = KernelSignature::create({ty}, out_ty);
                KCHECK_OK(func->add_kernel(kernel));
            }
        }

        const FunctionDoc is_in_doc{
                "Find each element in a set of values",
                ("For each element in `values`, return true if it is found in a given\n"
                 "set of values, false otherwise.\n"
                 "The set of values to look for must be given in SetLookupOptions.\n"
                 "By default, nulls are matched against the value set, this can be\n"
                 "changed in SetLookupOptions."),
                {"values"},
                "SetLookupOptions",
                /*options_required=*/true};

        const FunctionDoc is_in_meta_doc{
                "Find each element in a set of values",
                ("For each element in `values`, return true if it is found in `value_set`,\n"
                 "false otherwise."),
                {"values", "value_set"}};

        const FunctionDoc index_in_doc{
                "Return index of each element in a set of values",
                ("For each element in `values`, return its index in a given set of\n"
                 "values, or null if it is not found there.\n"
                 "The set of values to look for must be given in SetLookupOptions.\n"
                 "By default, nulls are matched against the value set, this can be\n"
                 "changed in SetLookupOptions."),
                {"values"},
                "SetLookupOptions",
                /*options_required=*/true};

        const FunctionDoc index_in_meta_doc{
                "Return index of each element in a set of values",
                ("For each element in `values`, return its index in the `value_set`,\n"
                 "or null if it is not found there."),
                {"values", "value_set"}};

// Enables calling is_in with call_function as though it were binary.
        class IsInMetaBinary : public MetaFunction {
        public:
            IsInMetaBinary() : MetaFunction("is_in_meta_binary", Arity::Binary(), is_in_meta_doc) {}

            turbo::Result<Datum> ExecuteImpl(const std::vector<Datum> &args,
                                             const FunctionOptions *options,
                                             ExecContext *ctx) const override {
                if (options != nullptr) {
                    return turbo::invalid_argument_error("Unexpected options for 'is_in_meta_binary' function");
                }
                return IsIn(args[0], args[1], ctx);
            }
        };

// Enables calling index_in with call_function as though it were binary.
        class IndexInMetaBinary : public MetaFunction {
        public:
            IndexInMetaBinary()
                    : MetaFunction("index_in_meta_binary", Arity::Binary(), index_in_meta_doc) {}

            turbo::Result<Datum> ExecuteImpl(const std::vector<Datum> &args,
                                             const FunctionOptions *options,
                                             ExecContext *ctx) const override {
                if (options != nullptr) {
                    return turbo::invalid_argument_error("Unexpected options for 'index_in_meta_binary' function");
                }
                return IndexIn(args[0], args[1], ctx);
            }
        };

        struct SetLookupFunction : ScalarFunction {
            using ScalarFunction::ScalarFunction;

            turbo::Result<const Kernel *> dispatch_best(std::vector<TypeHolder> *values) const override {
                EnsureDictionaryDecoded(values);
                return dispatch_exact(*values);
            }
        };

    }  // namespace

    void register_scalar_set_lookup(FunctionRegistry *registry) {
        // IsIn writes its boolean output into preallocated memory
        {
            ScalarKernel isin_base;
            isin_base.init = InitSetLookup;
            isin_base.exec = ExecIsIn;
            isin_base.null_handling = NullHandling::COMPUTED_PREALLOCATE;
            auto is_in = std::make_shared<SetLookupFunction>("is_in", Arity::Unary(), is_in_doc);

            AddBasicSetLookupKernels(isin_base, /*output_type=*/boolean(), is_in.get());

            isin_base.signature = KernelSignature::create({null()}, boolean());
            KCHECK_OK(is_in->add_kernel(isin_base));
            KCHECK_OK(registry->add_function(is_in));

            KCHECK_OK(registry->add_function(std::make_shared<IsInMetaBinary>()));
        }

        // IndexIn writes its int32 output into preallocated memory
        {
            ScalarKernel index_in_base;
            index_in_base.init = InitSetLookup;
            index_in_base.exec = ExecIndexIn;
            index_in_base.null_handling = NullHandling::COMPUTED_PREALLOCATE;
            auto index_in =
                    std::make_shared<SetLookupFunction>("index_in", Arity::Unary(), index_in_doc);

            AddBasicSetLookupKernels(index_in_base, /*output_type=*/int32(), index_in.get());

            index_in_base.signature = KernelSignature::create({null()}, int32());
            KCHECK_OK(index_in->add_kernel(index_in_base));
            KCHECK_OK(registry->add_function(index_in));

            KCHECK_OK(registry->add_function(std::make_shared<IndexInMetaBinary>()));
        }
    }

}  // namespace nebula::compute::internal
